#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

智能决策引擎模块
提供基于系统环境和检测结果的智能推荐功能
"""

import os
import sys
import json
import time
from collections import defaultdict

try:
    from core.utils import color_red, color_green, color_yellow, color_blue, color_purple
except ImportError:
    def color_red(text): return text
    def color_green(text): return text
    def color_yellow(text): return text
    def color_blue(text): return text
    def color_purple(text): return text

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()


class IntelligentDecisionEngine:
    """智能决策引擎"""
    
    def __init__(self):
        """初始化智能决策引擎"""
        self.system_info = {}
        self.check_results = {}
        self.backdoor_conditions = {}
        self.recommendations = []
        self.risk_score = 0
        self.confidence_level = 0.0
        
        # 风险评估权重
        self.risk_weights = {
            'ssh_access': 30,
            'root_access': 25,
            'network_exposure': 20,
            'service_running': 15,
            'file_permissions': 10
        }
        
        # 后门类型优先级
        self.backdoor_priorities = {
            'ssh_key': {'priority': 1, 'stealth': 9, 'persistence': 10},
            'crontab': {'priority': 2, 'stealth': 7, 'persistence': 9},
            'bash_alias': {'priority': 3, 'stealth': 8, 'persistence': 6},
            'strace': {'priority': 4, 'stealth': 6, 'persistence': 7},
            'system_service': {'priority': 5, 'stealth': 5, 'persistence': 10},
            'environment': {'priority': 6, 'stealth': 8, 'persistence': 5},
            'ssh_config': {'priority': 7, 'stealth': 7, 'persistence': 8},
            'network': {'priority': 8, 'stealth': 4, 'persistence': 9},
            'user_account': {'priority': 9, 'stealth': 3, 'persistence': 10}
        }
    
    def analyze_system_environment(self, system_info):
        """分析系统环境"""
        color_blue("[*] 分析系统环境...")
        self.system_info = system_info
        
        # 分析操作系统类型
        os_type = system_info.get('os_type', 'unknown')
        distro = system_info.get('distro', 'unknown')
        
        color_green("[+] 检测到系统: {} ({})".format(distro, os_type))
        
        # 分析系统版本和特性
        self._analyze_os_features(os_type, distro)
        
        return True
    
    def _analyze_os_features(self, os_type, distro):
        """分析操作系统特性"""
        features = {
            'systemd_support': False,
            'cron_support': True,
            'ssh_support': True,
            'package_manager': 'unknown'
        }
        
        # 基于发行版设置特性
        if 'ubuntu' in distro.lower() or 'debian' in distro.lower():
            features.update({
                'systemd_support': True,
                'package_manager': 'apt'
            })
        elif 'centos' in distro.lower() or 'rhel' in distro.lower() or 'fedora' in distro.lower():
            features.update({
                'systemd_support': True,
                'package_manager': 'yum'
            })
        elif 'alpine' in distro.lower():
            features.update({
                'systemd_support': False,
                'package_manager': 'apk'
            })
        
        self.system_info['features'] = features
        color_yellow("[?] 系统特性: {}".format(features))
    
    def evaluate_check_results(self, check_results):
        """评估检测结果"""
        color_blue("[*] 评估检测结果...")
        self.check_results = check_results
        
        # 计算风险评分
        self._calculate_risk_score()
        
        # 分析后门可行性
        self._analyze_backdoor_feasibility()
        
        # 生成推荐方案
        self._generate_recommendations()
        
        color_green("[+] 风险评分: {}/100".format(self.risk_score))
        color_green("[+] 置信度: {:.2f}".format(self.confidence_level))
        
        return True
    
    def _calculate_risk_score(self):
        """计算风险评分"""
        score = 0
        
        # SSH访问风险
        if self.check_results.get('ssh_user_check', False):
            score += self.risk_weights['ssh_access']
        
        # Root权限风险
        if self.check_results.get('root_access', False):
            score += self.risk_weights['root_access']
        
        # 网络暴露风险
        if self.check_results.get('network_exposure', False):
            score += self.risk_weights['network_exposure']
        
        # 服务运行风险
        if self.check_results.get('services_running', False):
            score += self.risk_weights['service_running']
        
        # 文件权限风险
        if self.check_results.get('file_permissions', False):
            score += self.risk_weights['file_permissions']
        
        self.risk_score = min(score, 100)
        self.confidence_level = self.risk_score / 100.0
    
    def _analyze_backdoor_feasibility(self):
        """分析后门可行性"""
        color_blue("[*] 分析后门可行性...")
        
        for backdoor_type, properties in self.backdoor_priorities.items():
            feasible = self._check_backdoor_conditions(backdoor_type)
            self.backdoor_conditions[backdoor_type] = {
                'feasible': feasible,
                'priority': properties['priority'],
                'stealth': properties['stealth'],
                'persistence': properties['persistence'],
                'score': self._calculate_backdoor_score(backdoor_type, feasible)
            }
    
    def _check_backdoor_conditions(self, backdoor_type):
        """检查特定后门的条件"""
        conditions = {
            'ssh_key': self.check_results.get('ssh_user_check', False),
            'crontab': self.check_results.get('cron_available', True),
            'bash_alias': self.check_results.get('bash_available', True),
            'strace': self.check_results.get('strace_available', False),
            'system_service': self.system_info.get('features', {}).get('systemd_support', False),
            'environment': True,  # 环境变量后门通常都可行
            'ssh_config': self.check_results.get('ssh_config_writable', False),
            'network': self.check_results.get('network_access', True),
            'user_account': self.check_results.get('user_creation_possible', False)
        }
        
        return conditions.get(backdoor_type, False)
    
    def _calculate_backdoor_score(self, backdoor_type, feasible):
        """计算后门评分"""
        if not feasible:
            return 0
        
        properties = self.backdoor_priorities[backdoor_type]
        
        # 基础分数 = (隐蔽性 + 持久性) / 2
        base_score = (properties['stealth'] + properties['persistence']) / 2
        
        # 优先级调整 (优先级越低数字越小，分数越高)
        priority_bonus = (10 - properties['priority']) * 0.5
        
        # 系统适配性调整
        compatibility_bonus = self._get_compatibility_bonus(backdoor_type)
        
        total_score = base_score + priority_bonus + compatibility_bonus
        return min(total_score, 10.0)
    
    def _get_compatibility_bonus(self, backdoor_type):
        """获取系统兼容性加分"""
        bonus = 0
        features = self.system_info.get('features', {})
        
        if backdoor_type == 'system_service' and features.get('systemd_support'):
            bonus += 1.0
        elif backdoor_type == 'ssh_key' and features.get('ssh_support'):
            bonus += 0.5
        elif backdoor_type == 'crontab' and features.get('cron_support'):
            bonus += 0.5
        
        return bonus
    
    def _generate_recommendations(self):
        """生成推荐方案"""
        color_blue("[*] 生成智能推荐方案...")
        
        # 按评分排序后门
        sorted_backdoors = sorted(
            self.backdoor_conditions.items(),
            key=lambda x: x[1]['score'],
            reverse=True
        )
        
        # 生成推荐
        self.recommendations = []
        
        for backdoor_type, info in sorted_backdoors:
            if info['feasible'] and info['score'] > 5.0:
                recommendation = {
                    'type': backdoor_type,
                    'score': info['score'],
                    'priority': info['priority'],
                    'stealth': info['stealth'],
                    'persistence': info['persistence'],
                    'reason': self._get_recommendation_reason(backdoor_type, info)
                }
                self.recommendations.append(recommendation)
        
        # 限制推荐数量
        self.recommendations = self.recommendations[:5]
    
    def _get_recommendation_reason(self, backdoor_type, info):
        """获取推荐理由"""
        reasons = {
            'ssh_key': "SSH密钥后门具有极高的隐蔽性和持久性，是首选方案",
            'crontab': "定时任务后门易于部署且持久性强，适合长期维持",
            'bash_alias': "Bash别名后门隐蔽性好，对系统影响小",
            'strace': "Strace后门可监控系统调用，具有强大的信息收集能力",
            'system_service': "系统服务后门持久性最强，但可能被检测",
            'environment': "环境变量后门简单易用，适合快速部署",
            'ssh_config': "SSH配置后门可控制SSH行为，隐蔽性较好",
            'network': "网络后门可建立持久连接，但网络特征明显",
            'user_account': "用户账户后门持久性强，但容易被发现"
        }
        
        base_reason = reasons.get(backdoor_type, "推荐使用此后门类型")
        score_reason = "评分: {:.1f}/10.0".format(info['score'])
        
        return "{} ({})".format(base_reason, score_reason)
    
    def get_intelligent_recommendations(self):
        """获取智能推荐"""
        if not self.recommendations:
            return []
        
        color_purple("\n" + "="*60)
        color_purple("[*] 🧠 智能推荐方案")
        color_purple("="*60)
        
        for i, rec in enumerate(self.recommendations, 1):
            color_green("\n[{}] {} 后门".format(i, rec['type'].upper()))
            color_yellow("    评分: {:.1f}/10.0".format(rec['score']))
            color_yellow("    隐蔽性: {}/10".format(rec['stealth']))
            color_yellow("    持久性: {}/10".format(rec['persistence']))
            color_blue("    推荐理由: {}".format(rec['reason']))
        
        return self.recommendations
    
    def get_deployment_strategy(self):
        """获取部署策略"""
        if not self.recommendations:
            return None
        
        strategy = {
            'primary': self.recommendations[0] if self.recommendations else None,
            'backup': self.recommendations[1:3] if len(self.recommendations) > 1 else [],
            'risk_level': self._assess_risk_level(),
            'stealth_mode': self.confidence_level > 0.7,
            'multi_layer': len(self.recommendations) >= 3
        }
        
        color_purple("\n" + "="*60)
        color_purple("[*] 🎯 部署策略建议")
        color_purple("="*60)
        
        if strategy['primary']:
            color_green("\n[主要方案] {}".format(strategy['primary']['type'].upper()))
            color_yellow("  推荐理由: {}".format(strategy['primary']['reason']))
        
        if strategy['backup']:
            color_blue("\n[备用方案]")
            for backup in strategy['backup']:
                color_blue("  - {}".format(backup['type'].upper()))
        
        color_yellow("\n[风险等级] {}".format(strategy['risk_level']))
        color_yellow("[隐蔽模式] {}".format('启用' if strategy['stealth_mode'] else '禁用'))
        color_yellow("[多层部署] {}".format('推荐' if strategy['multi_layer'] else '不推荐'))
        
        return strategy
    
    def _assess_risk_level(self):
        """评估风险等级"""
        if self.risk_score >= 80:
            return "极高"
        elif self.risk_score >= 60:
            return "高"
        elif self.risk_score >= 40:
            return "中等"
        elif self.risk_score >= 20:
            return "低"
        else:
            return "极低"
    
    def optimize_parameters(self, backdoor_type):
        """优化后门参数"""
        color_blue("[*] 优化 {} 后门参数...".format(backdoor_type))
        
        optimizations = {
            'ssh_key': {
                'key_type': 'ed25519',  # 更安全的密钥类型
                'key_comment': 'system-update',  # 伪装注释
                'file_permissions': '600',
                'stealth_tips': ['使用系统相关的密钥注释', '设置合适的文件权限', '避免在明显位置存放']
            },
            'crontab': {
                'schedule': '*/30 * * * *',  # 每30分钟执行
                'command_disguise': '/usr/bin/system-check',
                'log_redirect': '>/dev/null 2>&1',
                'stealth_tips': ['使用系统维护相关的命令名', '重定向输出避免日志', '选择合适的执行频率']
            },
            'bash_alias': {
                'alias_name': 'ls',  # 伪装成常用命令
                'original_backup': True,
                'stealth_tips': ['选择常用命令作为别名', '备份原始命令', '保持命令行为一致性']
            },
            'system_service': {
                'service_name': 'system-monitor',
                'description': 'System Monitoring Service',
                'start_type': 'auto',
                'stealth_tips': ['使用系统相关的服务名', '设置合理的服务描述', '配置自动启动']
            }
        }
        
        params = optimizations.get(backdoor_type, {})
        
        if params:
            color_green("[+] {} 优化参数:".format(backdoor_type))
            for key, value in params.items():
                if key != 'stealth_tips':
                    color_yellow("    {}: {}".format(key, value))
            
            if 'stealth_tips' in params:
                color_blue("  [隐蔽性建议]")
                for tip in params['stealth_tips']:
                    color_blue("    - {}".format(tip))
        
        return params
    
    def generate_execution_plan(self):
        """生成执行计划"""
        if not self.recommendations:
            return None
        
        plan = {
            'phases': [],
            'estimated_time': 0,
            'risk_mitigation': [],
            'rollback_plan': []
        }
        
        # 阶段1: 环境准备
        plan['phases'].append({
            'phase': 1,
            'name': '环境准备',
            'tasks': ['检查系统权限', '验证依赖工具', '创建工作目录'],
            'duration': 2
        })
        
        # 阶段2: 主要后门部署
        if self.recommendations:
            primary = self.recommendations[0]
            plan['phases'].append({
                'phase': 2,
                'name': "部署{0}后门".format(primary["type"]),
                'tasks': ["配置{0}后门".format(primary["type"]), '测试后门功能', '验证隐蔽性'],
                'duration': 5
            })
        
        # 阶段3: 备用后门部署
        if len(self.recommendations) > 1:
            plan['phases'].append({
                'phase': 3,
                'name': '部署备用后门',
                'tasks': ['部署备用后门', '配置冗余机制', '测试多层访问'],
                'duration': 8
            })
        
        # 阶段4: 清理和验证
        plan['phases'].append({
            'phase': 4,
            'name': '清理和验证',
            'tasks': ['清理部署痕迹', '验证后门功能', '生成访问文档'],
            'duration': 3
        })
        
        plan['estimated_time'] = sum(phase['duration'] for phase in plan['phases'])
        
        # 风险缓解措施
        plan['risk_mitigation'] = [
            '定期检查后门状态',
            '监控系统日志异常',
            '准备快速清理脚本',
            '建立安全通信渠道'
        ]
        
        # 回滚计划
        plan['rollback_plan'] = [
            '停止所有后门进程',
            '删除后门文件',
            '恢复原始配置',
            '清理系统日志'
        ]
        
        color_purple("\n" + "="*60)
        color_purple("[*] 📋 执行计划")
        color_purple("="*60)
        
        for phase in plan['phases']:
            color_green("\n阶段 {0}: {1} (预计 {2} 分钟)".format(phase['phase'], phase['name'], phase['duration']))
            for task in phase['tasks']:
                color_yellow("  - {0}".format(task))
        
        color_blue("\n总预计时间: {0} 分钟".format(plan['estimated_time']))
        
        return plan


def test_intelligent_decision():
    """测试智能决策引擎"""
    print("="*60)
    color_blue("[*] 智能决策引擎测试")
    print("="*60)
    
    # 创建决策引擎
    engine = IntelligentDecisionEngine()
    
    # 模拟系统信息
    system_info = {
        'os_type': 'Linux',
        'distro': 'Ubuntu 20.04',
        'kernel': '5.4.0',
        'arch': 'x86_64'
    }
    
    # 模拟检测结果
    check_results = {
        'ssh_user_check': True,
        'root_access': False,
        'network_exposure': True,
        'services_running': True,
        'file_permissions': True,
        'cron_available': True,
        'bash_available': True,
        'strace_available': False,
        'ssh_config_writable': True,
        'network_access': True,
        'user_creation_possible': False
    }
    
    # 执行分析
    engine.analyze_system_environment(system_info)
    engine.evaluate_check_results(check_results)
    
    # 获取推荐
    recommendations = engine.get_intelligent_recommendations()
    strategy = engine.get_deployment_strategy()
    plan = engine.generate_execution_plan()
    
    # 优化参数示例
    if recommendations:
        engine.optimize_parameters(recommendations[0]['type'])
    
    color_green("\n[+] 智能决策引擎测试完成!")


if __name__ == "__main__":
    test_intelligent_decision()